import mmap, time, ast, struct

###
# init
###

response.title = 'Firewall WebUI'



###
# constants
###

protocol = {
        1  : 'ICMP',
        6  : 'TCP',
        17 : 'UDP',
        255: 'RES',
    }
for i in range(256):
    if not protocol.get(i): protocol[i] = i


reason = {
        -1  : 'firewall inactive',
        -2  : 'not IPv4',
        -3  : 'prot not enforced',
        -4  : 'no matching rule',
        -5  : 'out of state',
        -6  : 'conn tab full',
        -7  : 'xmas packet',
    }

modes = {
        'active':          0x01,
        'icmp':            0x02,
        'tcp':             0x04,
        'udp':             0x08,
        'conn_track':      0x10,
        'cleanup_accept':  0x20,
    }

action = {
        0: 'DROP',
        1: 'ACCEPT',
    }

hook = {
    i: s for i, s in enumerate(['PRE', 'IN', 'FW', 'OUT', 'POST'])
    }



structs = {
    'log': 'I 3B 2I 2H i I',
    'rule': '>4B 2H 2I',
    'conn': '2I 2H I B 3x',
}


table_colors = {
    'row_even': '#DBE9FF',
    'row_odd': '#D8DADF',
    'title': '#A4AECE',
    }


conn_stat = {
    1:  'Syn sent',
    2:  'Syn ack',
    3:  'Estab',
    4:  'Closing',
}


###
# helpers
###

def int2ip(n):
    return '{}.{}.{}.{}'.format(
        (n >> 24) & 0xFF, (n >> 16) & 0xFF,
        (n >> 8) & 0xFF, n & 0xFF)

def ip2int(s):
    four_bytes = struct.pack('4B', *[int(i) for i in s.split('.')])
    return struct.unpack('>I', four_bytes)[0]


###
# configuration
###


class Config:
    def config(self, opt = None):
        with open('/sys/class/fw4/fw4_rules/config') as f:
            e = ast.literal_eval(f.readline())
        if opt:
            e ^= modes[opt]
            with open('/sys/class/fw4/fw4_rules/config', 'w') as f:
                f.write(str(e))

        return {k: (True if e & modes[k] else False) for k in modes.keys()}

    def make_config_button(self, opt):
        e = self.config()
        return A(
            'ON' if e[opt] else 'OFF',
            _href = URL('conf', 'config', args = opt),
            cid = 'config_button_{}'.format(opt),
            _class = 'button',
            _style = 'color: {}'.format('green' if e[opt] else 'red'))




###
# log
###



class Log:
    def __init__(self, l):
        if not isinstance(l, str) or len(l) != struct.calcsize(structs['log']):
            raise Exception('bad log init: l = {}'.format(l))

        (self.time, self.prot, self.action, self.hooknum,
            self.src_ip, self.dst_ip,
            self.src_port, self.dst_port,
            self.reason, self.count) = struct.unpack(structs['log'], l)


    def make_row(self, i):

        return TR(TD(str(i)),
            TD(time.ctime(self.time)), TD(protocol[self.prot]), TD(action[self.action]),
            TD(hook[self.hooknum]), TD(int2ip(self.src_ip)), TD(str(self.src_port) if protocol[self.prot] != 'ICMP' else ''),
            TD(int2ip(self.dst_ip)), TD(str(self.dst_port) if protocol[self.prot] != 'ICMP' else ''), TD(str(self.count)),
            TD(reason.get(self.reason) or ('rule #' + str(self.reason))),
            _style = 'background-color: {}'.format(table_colors['row_even'] if i % 2 == 0 else table_colors['row_odd']))

class Logs:
    def __init__(self):
        self.logs = []
        with open('/dev/fw4_log', 'rb') as log, \
             open('/sys/class/fw4/fw4_log/log_size') as log_size_file:

             log_size = ast.literal_eval(log_size_file.readline())
             l = mmap.mmap(log.fileno(), log_size, mmap.MAP_SHARED, mmap.PROT_READ)

             i = 0
             while ord(l[i * 28 + 4]):
                 self.logs.append(Log(l[i * 28: (i + 1) * 28]))
                 i += 1

             l.close

    def get_logs_as_table(self, f, t):
        return TABLE(
            TR(TD('#'), TD('Time'), TD('Protocol'), TD('Action'), TD('Hook'), TD('Src IP'),
            TD('Src Port'), TD('Dest IP'), TD('Dest Port'), TD('Count'),
            TD('Reason'), _style = 'background-color: {}'.format(table_colors['title'])),
            *[l.make_row(i + f + 1) for i, l in enumerate(self.logs[f: t])])




###
# rules
###

class Rules:
    def __init__(self):
        self.rules = []
        with open('/dev/fw4_rules', 'r+b') as rule_base, \
            open('/sys/class/fw4/fw4_rules/rules_size') as rule_base_size_file:

            rule_base_size = ast.literal_eval(rule_base_size_file.readline())
            r = mmap.mmap(rule_base.fileno(), rule_base_size)
            i = 0
            while ord(r[i * 16]) != 255:
                self.rules.append(Rule(r[i * 16: (i + 1) * 16]))
                i += 1

            r.close()

    def get_as_table(self):
        return DIV(TABLE(
            TR(TD('#'), TD('Protocol'), TD('Src IP'), TD('Src Mask'), TD('Src Port'),
            TD('Dest IP'), TD('Dest Mask'), TD('Dest Port'), TD('Action'),
            TD('Edit', _id = 'rule_edit'),
            _style = 'background-color: {}'.format(table_colors['title'])),
            *[r.make_row(i) for i, r in enumerate(self.rules)]),
            )

    def commit(self):
        with open('/dev/fw4_rules', 'r+b') as rule_base, \
            open('/sys/class/fw4/fw4_rules/rules_size') as rule_base_size_file:

            rule_base_size = ast.literal_eval(rule_base_size_file.readline())
            m = mmap.mmap(rule_base.fileno(), rule_base_size)

            for i, r in enumerate(self.rules + [Rule()]):
                m[i * 16: (i + 1) * 16] = r.make_bytes()

            m.close()

class Rule:
    def __init__(self, r = None, form = None):
        
        if form:
            self.protocol = int(form.vars.protocol)
            self.src_ip   = ip2int(form.vars.src_ip)
            self.src_mask = form.vars.src_mask
            self.src_port = form.vars.src_port
            self.dst_ip   = ip2int(form.vars.dst_ip)
            self.dst_mask = form.vars.dst_mask
            self.dst_port = form.vars.dst_port
            self.action   = int(form.vars.action)
        elif r:
            (self.protocol, self.src_mask, self.dst_mask, self.action,
                self.src_port, self.dst_port, self.src_ip,
                self.dst_ip) = struct.unpack(structs['rule'], r)
        else:
            self.protocol = 255
            self.src_ip, self.src_mask, self.src_port, self.dst_ip = 0, 0, 0, 0
            self.dst_mask, self.dst_port, self.action = 0, 0, 0

    def make_row(self, i):
        return TR(
            TD(str(i)), TD(protocol[self.protocol]), TD(int2ip(self.src_ip)), TD(str(self.src_mask)),
            TD(str(self.src_port) if ((protocol[self.protocol] in ['TCP', 'UDP']) and self.src_port != 0) else ''),
            TD(int2ip(self.dst_ip)), TD(str(self.dst_mask)),
            TD(str(self.dst_port) if ((protocol[self.protocol] in ['TCP', 'UDP']) and self.dst_port != 0) else ''),
            TD(str(action[self.action])),
            TD(self.make_crud(i), _id = 'rule_edit_{}'.format(i)),
            _style = 'background-color: {}'.format(table_colors['row_even'] if i % 2 == 0 else table_colors['row_odd']),
            _id = 'rule_{}'.format(i),
            )

    def make_bytes(self):
        return struct.pack(structs['rule'],
            self.protocol, self.src_mask, self.dst_mask, self.action,
            self.src_port, self.dst_port, self.src_ip, self.dst_ip)


    def make_crud(self, i):
        return DIV(*[A(s[:3], _href = URL('rule', s, args = i), cid = 'rules_display', _style = 'padding: 0px 5px') \
                     for s in ('delete', 'add', 'duplicate', 'edit', 'up', 'down')])

    def make_form(self):
        return SQLFORM.factory(
            Field('protocol', requires = IS_IN_SET(protocol, zero = None), default = self.protocol),
            Field('src_ip', requires = IS_IPV4(), default = int2ip(self.src_ip)),
            Field('src_mask', requires = IS_INT_IN_RANGE(0, 33), default = self.src_mask),
            Field('src_port', requires = IS_INT_IN_RANGE(0, 1 << 16), default = self.src_port),
            Field('dst_ip', requires = IS_IPV4(), default = int2ip(self.dst_ip)),
            Field('dst_mask', requires = IS_INT_IN_RANGE(0, 33), default = self.dst_mask),
            Field('dst_port', requires = IS_INT_IN_RANGE(0, 1 << 16), default = self.dst_port),
            Field('action', requires = IS_IN_SET(action, zero = None), default = self.action),
            )


###
# Connection table
###


class ConnTabEntry:
    def __init__(self, c):
        if not isinstance(c, str) or len(c) != struct.calcsize(structs['conn']):
            raise Exception('bad conn init: c = {}'.format(c))

        (self.cli_ip, self.ser_ip, self.cli_port, self.ser_port,
            self.expires, self.state) = struct.unpack(structs['conn'], c)


    def make_row(self, i):

        return TR(TD(str(i)),
            TD(int2ip(self.cli_ip)), TD(str(self.cli_port)),
            TD(int2ip(self.ser_ip)), TD(str(self.ser_port)),
            TD(time.ctime(self.expires)), TD(conn_stat[self.state]),
            
            _style = 'background-color: {}'.format(table_colors['row_even'] if i % 2 == 0 else table_colors['row_odd']))

class ConnTab:
    def __init__(self):
        self.entries = []
        with open('/dev/fw4_conn_tab', 'rb') as conn_tab, \
             open('/sys/class/fw4/fw4_conn_tab/conn_tab_size') as conn_tab_size_file:
             
             ent_size = struct.calcsize(structs['conn'])
             now = time.time()

             conn_tab_size = ast.literal_eval(conn_tab_size_file.readline())
             m = mmap.mmap(conn_tab.fileno(), conn_tab_size, mmap.MAP_SHARED, mmap.PROT_READ)

             for i in range(0, conn_tab_size, ent_size):
                 ent = ConnTabEntry(m[i: i + ent_size])
                 if ent.expires >= now:
                     self.entries.append(ent)

             m.close

    def get_as_table(self):
        return TABLE(
            TR(TD('#'), TD('Cli IP'), TD('Cli port'), TD('Ser IP'), TD('Ser port'), 
            TD('Expires'), TD('State'),
            
            _style = 'background-color: {}'.format(table_colors['title'])),
            *[e.make_row(i) for i, e in enumerate(self.entries)])
